import torch
import torchvision
import torch.nn as nn
import torch.optim as optim
from torch.nn import functional as F
from utils import one_hot

# config data
batch_size = 512
# step1:load data
train_loader = torch.utils.data.DataLoader(
    torchvision.datasets.MNIST('mnist_data', train=True, download=True,
                               transform=torchvision.transforms.Compose([
                                   torchvision.transforms.ToTensor(),
                                   torchvision.transforms.Normalize(
                                       (0.1307,), (0.3081,))
                               ])),
    batch_size=batch_size, shuffle=True)

test_loader = torch.utils.data.DataLoader(
    torchvision.datasets.MNIST('mnist_data/', train=False, download=True,
                               transform=torchvision.transforms.Compose([
                                   torchvision.transforms.ToTensor(),
                                   torchvision.transforms.Normalize(
                                       (0.1307,), (0.3081,))
                               ])),
    batch_size=batch_size, shuffle=True)

# step2:build model


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        # loss function
        self.cretenon = F.mse_loss
        # model
        self.CNN_unit = nn.Sequential(
            nn.Conv2d(1, 6, kernel_size=5, stride=1, padding=1),
            nn.ReLU(),
            nn.MaxPool2d((2, 2)),
            nn.Conv2d(6, 16, kernel_size=5, stride=1, padding=1),
            nn.ReLU(),
            nn.MaxPool2d((2, 2)),
        )
        self.DNN_unit = nn.Sequential(
            nn.Linear(400, 256),
            nn.ReLU(),
            nn.Linear(256, 64),
            nn.ReLU(),
            nn.Linear(64, 10),
        )

    def forward(self, x):
        x = self.CNN_unit(x)
        # [b,12,28,28]
        x = x.view(x.size(0), -1)
        x = self.DNN_unit(x)
        return x


# step3:train
# model
net = Net()
# Optimizer
optimizer = optim.Adam(net.parameters(), lr=1e-3)


def main():
    for epoch in range(3):
        for batch_idx, (x, y) in enumerate(train_loader):
            out = net(x)
            # [b,10]
            y_onehot = one_hot(y)
            # loss = mse(out,y_onehot)
            loss = net.cretenon(out, y_onehot)
            optimizer.zero_grad()
            loss.backward()
            # w' = w - lr * grad
            optimizer.step()
            if batch_idx % 10 == 0:
                print(epoch, batch_idx, loss.item())

    # get optimal [w1 w2 w3 b1 b2 b3]
    total_correct = 0
    for x, y in test_loader:
        out = net(x)
        pred = out.argmax(dim=1)
        correct = pred.eq(y).sum().float()
        total_correct += correct

    total_sum = len(test_loader.dataset)
    acc = total_correct / total_sum
    print('acc=', acc)


if __name__ == "__main__":
    main()
